# 切小小图，切很多份

import os, glob, shutil
import xml.dom.minidom as minidom
import xml.etree.ElementTree as ET
import cv2
from xml.dom.minidom import Document
import numpy as np
import fire
import sys
from numpy.random import randint
from utils.conve import *



def generate_xml(floder_data, filename_data, img_path, img_width, img_height, data_defect, save_path):
    '''
    写xml
    floder_data:xml对应的图片的相对于根目录的目录
    filename_data:对应的图片名称
    img_path:对应的图片的绝对路径
    img_width:图片宽
    img_hight:图片高
    data_defect:缺陷名称
    save_path:保存路径
    '''
    doc = Document()
    annotation_node = doc.createElement("annotation")
    doc.appendChild(annotation_node)

    #在annotation节点下添加floder文本节点
    floder_node = doc.createElement("floder")
    annotation_node.appendChild(floder_node)
    floder_node_data = doc.createTextNode(floder_data)#文件夹名
    floder_node.appendChild(floder_node_data)

    #在annotation节点下添加filename文本节点
    filename_node = doc.createElement("filename")
    annotation_node.appendChild(filename_node)
    filename_node_data = doc.createTextNode(filename_data)#文件名
    filename_node.appendChild(filename_node_data)

    #在annotation节点下添加path文本节点
    path_node = doc.createElement("path")
    annotation_node.appendChild(path_node)
    path_node_data = doc.createTextNode(img_path)#路径名
    path_node.appendChild(path_node_data)

    #添加source节点
    source_node = doc.createElement("source")
    annotation_node.appendChild(source_node)

    #在source节点下添加datebase文本节点
    database_node = doc.createElement("database")
    source_node.appendChild(database_node)
    database_node_data = doc.createTextNode("Unknown")
    database_node.appendChild(database_node_data)

    #在annotation添加size节点
    size_node = doc.createElement("size")
    annotation_node.appendChild(size_node)

    #在size节点下添加width文本节点
    width_node = doc.createElement("width")
    size_node.appendChild(width_node)
    width_node_data = doc.createTextNode(img_width)#图像宽
    width_node.appendChild(width_node_data)

    #在size节点下添加height文本节点
    height_node = doc.createElement("height")
    size_node.appendChild(height_node)
    height_node_data = doc.createTextNode(img_height)#图像高
    height_node.appendChild(height_node_data)

    #在size节点下添加depth文本节点
    depth_node = doc.createElement("depth")
    size_node.appendChild(depth_node)
    depth_node_data = doc.createTextNode("1")#图像通道数
    depth_node.appendChild(depth_node_data)

    #在annotation节点下添加segmented文本节点
    segmented_node = doc.createElement("segmented")
    annotation_node.appendChild(segmented_node)
    segmented_node_data = doc.createTextNode("0")
    segmented_node.appendChild(segmented_node_data)

    #在annotation添加object节点
    object_node = doc.createElement("object")
    annotation_node.appendChild(object_node)

    #在object节点下添加name文本节点
    name_node = doc.createElement("name")
    object_node.appendChild(name_node)
    name_node_data = doc.createTextNode('wg_' + data_defect)#缺陷名称
    name_node.appendChild(name_node_data)

    #在object节点下添加pose文本节点
    pose_node = doc.createElement("pose")
    object_node.appendChild(pose_node)
    pose_node_data = doc.createTextNode("Unspecified")
    pose_node.appendChild(pose_node_data)

    #在object节点下添加truncated文本节点
    truncated_node = doc.createElement("truncated")
    object_node.appendChild(truncated_node)
    truncated_node_data = doc.createTextNode("0")
    truncated_node.appendChild(truncated_node_data)

    #在object节点下添加pose文本节点
    difficult_node = doc.createElement("difficult")
    object_node.appendChild(difficult_node)
    difficult_node_data = doc.createTextNode("0")
    difficult_node.appendChild(difficult_node_data)

    #在object节点下添加bndbox节点
    bndbox_node = doc.createElement("bndbox")
    object_node.appendChild(bndbox_node)

    #在bndbox节点下添加xmin文本节点
    xmin_node = doc.createElement("xmin")
    bndbox_node.appendChild(xmin_node)
    xmin_node_data = doc.createTextNode("10")
    xmin_node.appendChild(xmin_node_data)

    #在bndbox节点下添加ymin文本节点
    ymin_node = doc.createElement("ymin")
    bndbox_node.appendChild(ymin_node)
    ymin_node_data = doc.createTextNode("10")
    ymin_node.appendChild(ymin_node_data)

    #在bndbox节点下添加xmax文本节点
    xmax_node = doc.createElement("xmax")
    bndbox_node.appendChild(xmax_node)
    xmax_node_data = doc.createTextNode("20")
    xmax_node.appendChild(xmax_node_data)

    #在bndbox节点下添加ymax文本节点
    ymax_node = doc.createElement("ymax")
    bndbox_node.appendChild(ymax_node)
    ymax_node_data = doc.createTextNode("20")
    ymax_node.appendChild(ymax_node_data)
    
    
    xml_path = filename_data.split(".")[0] + ".xml"
    file_save_path = os.path.join(save_path, xml_path)
    #print("path: ", file_save_path)
    f = open(file_save_path, "w", encoding='utf-8')
    f.write(doc.toprettyxml(indent = "    "))
    f.close()



def grid_cut_row_lines(img, img_name, row_lines, save_path, top_value, down_value, left_value, right_value, conve):
    '''
    切小图，切线上下拓展，图片左右紧缩
    img:读取的图片
    img_name:图片名字
    row_lines:切线位置
    save_path:保存路径
    top_value:上方扩展
    down_value:下方扩展
    left_value:左边紧缩
    right_value:右边紧缩
    conve:前处理方法，在conve.py中
    '''
    if left_value == 0 and right_value == 0:
        img_cut = img[row_lines-top_value:row_lines+down_value, :, :]
    else:
        img_cut = img[row_lines-top_value:row_lines+down_value, left_value:-right_value, :]
    # print(img_cut.shape)
    
    img_cut = cv2.resize(img_cut, (224, 224)) # 选择在此调整大小
    
    if conve:
        for one_conve in conve.split(','):
            print('img file {} in {}'.format(img_name, one_conve))
            img_cut = eval(one_conve)(img_cut)
        print('<==========end==========>')
    
    cv2.imwrite(os.path.join(save_path, img_name), img_cut)



# glob.glob(os.path.join(img_path, '*.jpg'))
def get_cut_pic(img_path, img_save_path, xml_save_path, count_num, mode, preb, conve):
    '''
    切小小图，传递上下左右参数，循环n次生成n张
    img_path:图片的目录
    img_save_path:保存小图的目录
    xml_save_path:保存xml的目录
    count_num:循环次数，即生成份数
    mode:传入pian/chuan（片/串），控制上下左右和中间切线和宽度
    preb:传入bingpian/bubingpian（并/不并），控制xml标签名和图片名
    conve:前处理方法，在conve.py中
    '''
    files = glob.glob(os.path.join(img_path, '*.jpg'))
    for k, file in enumerate(files):
        print('<=========={}%==========>'.format((k+1)/len(files)*100))
        print('img file {} is processing'.format(os.path.basename(file)))
        img = cv2.imdecode(np.fromfile(file, dtype=np.uint8), 1)
        hh, ww = img.shape[:2]
        # print(img.shape)
        for i in range(count_num):
            if mode == 'chuan':
                top_value = int(np.random.randint(120, 150, 1))
                down_value = 270 - top_value
                left_value = int(np.random.randint(35, 40, 1))
                right_value = 75 - left_value
                # left_value = 0
                # right_value = 0
                cut_line = hh // 2
                width = 375
                # top_value = cut_line - 1
                # down_value = cut_line - 1
            elif mode == 'pian':
                # pad 太大，超出边界
                cut_line = hh // 2
                top_value = int(np.random.randint(130, 140, 1))
                down_value = 270 - top_value
                left_value = int(np.random.randint(20, 30, 1))
                right_value = 50 - left_value
                width = 750
            file_name = preb + '_' + os.path.basename(file).split('.')[0] + '_' + str(i) + '.jpg'

            try:
                grid_cut_row_lines(img, file_name, cut_line, img_save_path, top_value, down_value, left_value, right_value, conve)           
            except:
                print("img {} cut error happened!".format(file_name))
            
            try:
                generate_xml(img_path, file_name, os.path.join(img_path, file_name), str(width - left_value - right_value), str(top_value + down_value), preb, xml_save_path)
            except:    
                print("xml {} make error happened!".format(file_name))


def check_dir(dir_path):
    '''
    检测路径，如果有路径提示是否删除，无路径新建路径
    '''
    if os.path.exists(dir_path):
        print('!!!dir {} exists'.format(dir_path))
        comfirm = input("!!!remove dir, yes or mix? [y or m]: ")
        if comfirm == 'y':
            print('!!!{} have been removed!!!'.format(dir_path))
            shutil.rmtree(dir_path)
            os.makedirs(dir_path)
        elif comfirm != 'm':
            print('did not remove!')
            sys.exit(0)
    else:
        os.makedirs(dir_path)



def check_cls(cls):
    '''
    根据读取的文件夹名定义mode和preb
    '''
    if cls == 'pbp':
        mode = 'pian'
        preb = 'bingpian'
    elif cls == 'pbbp':
        mode = 'pian'
        preb = 'bubingpian'
    elif cls == 'cbp':
        mode = 'chuan'
        preb = 'bingpian'
    elif cls == 'cbbp':
        mode = 'chuan'
        preb = 'bubingpian'
    return mode, preb



def main(root_path, save_path, count_nums = [1, 1, 1, 1], conve = False, dir = True):
    '''
    遍历根目录，切小图写xml，保存n份
    root_path:包含几种类型的图片目录的根目录
    save_path:最后保存小小图的目录
    count_nums:对应pbp、pbbp、cbp、cbbp（片并片、片不并片、串并片、串不并片）的生成份数
    conve:选择前处理方式，False为关闭，字符串为调用方法，以','隔开先后调用，具体方法在conve.py中
    dir:是否新建路径，True新建路径
    '''
    img_cls_path = os.listdir(root_path)
    img_save_path = os.path.join(save_path, 'img')
    xml_save_path = os.path.join(save_path, 'xml')
    if dir:
        check_dir(img_save_path)
        check_dir(xml_save_path)
    count_count = {}
    i = 0
    for cls in ['pbp', 'pbbp', 'cbp', 'cbbp']:
        count_count[cls] = count_nums[i]
        i += 1
    for img_fin_path in img_cls_path:
        print('<================dir:{}===============>'.format(os.path.dirname(img_fin_path)))
        img_path = os.path.join(root_path, img_fin_path)
        mode, preb = check_cls(img_fin_path)
        get_cut_pic(img_path, img_save_path, xml_save_path, count_count[img_fin_path], mode, preb, conve)



if __name__ == '__main__':
    fire.Fire(main)